from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
from nnunetv2.training.loss.compound_losses import DC_and_CE_loss, DC_and_BCE_loss, DC_SkelREC_and_DropCE_loss
from nnunetv2.training.loss.dice import MemoryEfficientSoftDiceLoss
from nnunetv2.training.loss.deep_supervision import DeepSupervisionWrapper
from nnunetv2.training.nnUNetTrainer.variants.loss.nnUNetTrainerSkeletonRecall import nnUNetTrainerSkeletonRecall
import numpy as np
import torch
from nnunetv2.utilities.helpers import softmax_helper_dim1
import warnings


class nnUNetTrainerDropCESkeletonRecall(nnUNetTrainerSkeletonRecall):
    def _build_loss(self):
        if self.label_manager.ignore_label is not None:
            warnings.warn('Support for ignore label with Skeleton Recall is experimental and may not work as expected')
        loss = DC_SkelREC_and_DropCE_loss(soft_dice_kwargs={'batch_dice': self.configuration_manager.batch_dice, 
                                                        'smooth': 1e-5, 'do_bg': False, 'ddp': self.is_ddp}, 
                                      soft_skelrec_kwargs={'batch_dice': self.configuration_manager.batch_dice, 
                                                           'smooth': 1e-5, 'do_bg': False, 'ddp': self.is_ddp}, 
                                      ce_kwargs={}, weight_ce=1, weight_dice=1, weight_srec=self.weight_srec, 
                                      ignore_label=self.label_manager.ignore_label, dice_class=MemoryEfficientSoftDiceLoss)

        if self.enable_deep_supervision:
            deep_supervision_scales = self._get_deep_supervision_scales()

            # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases
            # this gives higher resolution outputs more weight in the loss
            weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))])
            weights[-1] = 0

            # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1
            weights = weights / weights.sum()
            # now wrap the loss
            loss = DeepSupervisionWrapper(loss, weights)
        return loss